[Python] コードを簡潔に書ける代入式を学ぶ

こんにちは。サービス開発室の武田です。

Pythonのコードを普段から書いているのですが、Python 3.8で追加された「代入式」のことをあまり知らなかったので勉強してみました。

代入式とは

代入式はPython 3.8で追加された機能です。「代入」と「式の評価」を1ステートメントで書ける機能です。:=という記号を使用し、「セイウチ演算子」と呼ばれることもあります。

たとえば、「リストの長さが0より大きければその要素数をプリントする」という処理を考えてみます。従来であれば素直に書くと次のようになります。

some_list = get_some_list()
if len(some_list) > 0:
  print(f"list length: {len(some_list)}")

しかし2回len()関数を呼び出しているのが気持ち悪いです。そこで次のように書き直すでしょう。

some_list = get_some_list()
n = len(some_list)
if n > 0:
  print(f"list length: {n}")

len()関数の呼び出しが1回になりました。その点ではよくなりましたが、コードが1行増えてしまいました。追加された「代入式」を使用すると次のように、簡潔に記述できます。

some_list = get_some_list()
if (n := len(some_list)) > 0:
  print(f"list length: {n}")

このほかにも、「正規表現のマッチング」や「リスト内包表記」などの具体例がドキュメントで示されています。参考URLから確認してみてください。

実際のコードに適用できる部分を探してみた

実際に携わっているプロジェクトのコードに適用できる部分を探してみました。呼び出している関数名などは適当に変えてあります。

具体例1：取得したオブジェクトがNoneでなければ処理をする

関数などの戻り値をチェックしてから使用するオーソドックスな使い方です。おそらくこれが一番多い使い方なのではないでしょうか。

before

line_items = context.get("line_items")
if line_items is not None:
    return line_items

after

if (line_items := context.get("line_items")) is not None:
    return line_items

具体例2：ロール一覧から test- で始まるものを絞り込んで取得

リスト内包表記のif節で使用できます。これによって何度も同じ参照や計算をする必要がなくなります。今回は辞書の参照なのでコストは高くありませんが、複雑な計算や関数呼び出しなどの場合はよりメリットが強くなります。

before

list_roles = iam.list_roles()
role_names = [r["RoleName"] for r in list_roles["Roles"] if r["RoleName"].startswith("test-")]

after

list_roles = iam.list_roles()
role_names = [name for r in list_roles["Roles"] if (name := r["RoleName"]).startswith("test-")]

条件が複数のパターン。

list_roles = iam.list_roles()
role_names = [name for r in list_roles["Roles"] if (name := r["RoleName"]).startswith("test-") or name.startswith("mock-")]

具体例3：正規表現にマッチしたオブジェクトを操作する

正規表現にマッチした文字列に対して処理をするのも一般的な処理でしょう。本質的には具体例1とやっていることは同じです。

before

match = re.search(r"test_(.*\.py)", line)
if match:
    files.append(match.group())

after

if (match := re.search(r"test_(.*\.py)", line)):
    files.append(match.group())

elifが必要なケースでは簡潔さが際立ちます。

if (match := re.search(r"test_(.*\.py)", line)):
    files.append(match.group())
elif (match := re.search(r"mock_(.*\.py)", line)):
    pass

代入式を使わない場合、同様の処理は次のように書く必要があります。

match = re.search(r"test_(.*\.py)", line)
if match:
    files.append(match.group())
else:
    match = re.search(r"mock_(.*\.py)", line)
    if match:
        pass

少し簡潔に書こうとすると次のようになりますが、match2は無駄になる可能性もありやや不満が残ります。

match1 = re.search(r"test_(.*\.py)", line)
match2 = re.search(r"mock_(.*\.py)", line)
if match1:
    files.append(match.group())
elif match2:
    pass

具体例4：NextTokenが存在すればそれを使用して操作する

最後はAWSのList系操作でよくあるNextTokenです。ただ個人的に、これの良し悪しの比較は難しいと思っていまして、ご意見あれば伺いたいです。

before

while "NextToken" in roots:
    roots = client.list_roots(NextToken=roots["NextToken"])

after

while (next_token := roots.get("NextToken")):
    roots = client.list_roots(NextToken=next_token)

まとめ

コードを簡潔に書けるようになる「代入式」または「セイウチ演算子」の紹介でした。便利な反面、多用しすぎると可読性を落とす可能性があります。ドキュメントには「セイウチ演算子の使用は、複雑さを減らしたり可読性を向上させる綺麗なケースに限るよう努めてください。」とありますので、気をつけましょう。

参考URL

• What's New In Python 3.8 — Python 3.12.1 ドキュメント
• PEP 572 – Assignment Expressions | peps.python.org